import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.append(project_root)
import pickle
import argparse
import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
import numpy as np

from utils import *

NUM_FRAMES = None
X_RANGE = None
Y_RANGE = None

NUM_FRAMES_CACHE = 500

lat0 = None
lon0 = None
h0 = None

all_data = []

def process_mat_file(mat_data, output_folder, file_name, x_range, y_range):
    global lat0
    global lon0
    global h0
    global all_data

    start_frame_llh = mat_data.get('StartFrameLLH', None)
    start_frame_prh = mat_data.get('StartFrameRPY', None)
    azi_map = mat_data.get('aziLinearlityMap', None)
    ele_map = mat_data.get('eleLinearlityMap', None)
    error_map = mat_data.get('errorMap', None)
    target_x = mat_data.get('targetX', None)
    target_y = mat_data.get('targetY', None)
    target_z = mat_data.get('targetZ', None)
    dn = mat_data.get('DN', None)

    if not all([azi_map is not None, ele_map is not None, error_map is not None, target_x is not None, 
                target_y is not None, target_z is not None, dn is not None, start_frame_llh is not None, start_frame_prh is not None]):
        raise ".mat文件中缺少数据"
    
    lat = start_frame_llh[0][0]  # 纬度
    lon = start_frame_llh[0][1]  # 精度
    h = start_frame_llh[0][2] / 1000  # 海拔高度

    pitch = start_frame_prh[0][0]  # 俯仰角（重力基准）
    roll = start_frame_prh[0][1]  # 滚转角（重力基准）
    heading = start_frame_prh[0][2]  # 与北方向夹角

    if lat0 == None:
        lat0 = lat
        lon0 = lon
        h0 = h

    X = target_x.flatten()
    Y = target_y.flatten()
    Z = target_z.flatten()
    DN = dn.flatten()
    azi = azi_map.flatten()
    ele = ele_map.flatten()
    err = error_map.flatten()

    frame_data = {
        'x': X,
        'y': Y,
        'z': Z,
        'dn': DN,
        'aziLinearlityMap': azi,
        'eleLinearlityMap': ele,
        'errorMap': err
    }

    frame_data = filter_z(frame_data)
    frame_data['x'], frame_data['y'], frame_data['z'] = local_to_global_coords(frame_data['x'], frame_data['y'], frame_data['z'], lat, lon, h, pitch, roll, heading, lat0, lon0, h0)
    
    all_data.append(frame_data)
    
    if len(all_data) > NUM_FRAMES_CACHE:
        all_data.pop(0)
    
    generate_enhanced_data(x_range, y_range, output_folder, lat, lon, h, pitch, roll, heading, file_name)

def generate_enhanced_data(x_range, y_range, output_folder, current_lat, current_lon, current_h, current_pitch, current_roll, current_heading, file_name):
    global all_data
    
    current_e, current_n, current_u = ecef_to_enu(
        *latlonh_to_ecef(current_lat, current_lon, current_h),
        lat0, lon0, h0
    )
    current_origin = np.array([current_e, current_n, current_u]).reshape(3, 1)
    
    R = rotation_matrix_from_euler(0, 0, current_heading)

    valid_frames = []
    for frame in reversed(all_data):
        frame = frame.copy()
        points_global = np.vstack((frame['x'], frame['y'], frame['z']))
        points_local = R.T.dot(points_global - current_origin)
        frame['x'], frame['y'], frame['z'] = points_local[0, :], points_local[1, :], points_local[2, :]

        frame = filter_data(frame, x_range, y_range, filter_noise=False)

        valid_frames.append({key: np.array(value) for key, value in frame.items()})

    total_valid = len(valid_frames)

    if total_valid >= NUM_FRAMES:
        stride = total_valid / NUM_FRAMES
        selected_indices = [int(round(i * stride)) for i in range(NUM_FRAMES)]
        selected_frames = [valid_frames[i] for i in selected_indices]
    else:
        selected_frames = valid_frames.copy()
        while len(selected_frames) < NUM_FRAMES:
            selected_frames.append({})

    plot_point_cloud(selected_frames, output_folder, file_name, x_range, y_range)

    output_file = os.path.join(output_folder, f'{file_name}.pkl')
    try:
        with open(output_file, 'wb') as f:
            pickle.dump({
                'data': selected_frames,
                'lat': current_lat, 
                'lon': current_lon, 
                'h': current_h, 
                'pitch': current_pitch, 
                'roll': current_roll, 
                'heading': current_heading
                }, f)
    except Exception as e:
        print(f"保存增强数据到 {output_file} 时出错: {e}")

def plot_point_cloud(frames, output_folder, file_name, x_range, y_range):
    x, y, z = [], [], []

    for frame in frames:
        if len(frame) > 0:
            frame = filter_data(frame, x_range, y_range, filter_noise=True)
            x.extend(frame['x'])
            y.extend(frame['y'])
            z.extend(frame['z'])
    
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(projection='3d')
    ax.scatter(x, y, z, s=1)
    
    ax.view_init(elev=90, azim=-90)
    ax.set_xlim(x_range)
    ax.set_ylim(y_range)

    os.makedirs(output_folder + '/img', exist_ok=True)
    plt.savefig(f"{output_folder}/img/{file_name}", dpi=100, facecolor='white', bbox_inches='tight', pad_inches=0)
    plt.close()

def main():
    global X_RANGE, Y_RANGE, NUM_FRAMES

    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--X_RANGE", type=str, required=True)
    parser.add_argument("--Y_RANGE", type=str, required=True)
    parser.add_argument("--NUM_FRAMES", type=str, required=True)
    args = parser.parse_args()
    
    X_RANGE = eval(args.X_RANGE)
    Y_RANGE = eval(args.Y_RANGE)
    NUM_FRAMES = eval(args.NUM_FRAMES) 

    input_folder = "./workspace/data/SAR"
    output_folder = "./workspace/image/SAR/input"
    
    os.makedirs(output_folder, exist_ok=True)

    all_files = []
    with os.scandir(input_folder) as entries:
        files = [entry for entry in entries if entry.is_file() and entry.name.endswith(".mat")]
        files.sort(key=lambda entry: natural_key(entry.name))
        
        for file in files:
            all_files.append(file.path)

    total_files = len(all_files)
    for index, file_path in enumerate(all_files, start=1):
        print(f"正在处理文件 ({index}/{total_files})")

        mat_data = read_mat_file(file_path)
        
        file_name = os.path.splitext(os.path.basename(file_path))[0]

        # output_file = os.path.join(output_folder, f'{file_name}.pkl')
        # if os.path.exists(output_file):
        #     continue

        if mat_data:
            process_mat_file(
                mat_data,
                output_folder=output_folder,
                file_name=file_name,
                x_range=X_RANGE,
                y_range=Y_RANGE
            )
    
if __name__ == "__main__":
    main()